#!/usr/bin/env python
from __future__ import print_function, division
import warnings
import numpy as np
import pandas as pd
from scipy import interpolate

xcol_, ycol_, zcol_, wcol_ = 0, 1, None, None

markers = ['o', 's', '^', 'v', '+', 'x', '*', 'd', '<', '>']
nmarker = len(markers)
imarker = 0

linestyles = ['-', '--', '-.', ':']
nls = len(linestyles)
ils = 0

icolor = 0

def loader(infile, type_='table', xcol=0, ycol=1, zcol=None, wcol=None):
    global xcol_, ycol_, zcol_, wcol_

    try:
        xcol_ = int(xcol)
    except:
        xcol_ = xcol

    try:
        ycol_ = int(ycol)
    except:
        ycol_ = ycol

    try:
        zcol_ = int(zcol)
    except:
        zcol_ = zcol

    try:
        wcol_ = int(wcol)
    except:
        wcol_ = wcol

    if type_ == 'table':
        #dat = pd.read_table(infile, comment='#', delim_whitespace=True, header=None)
        dat = pd.read_csv(infile, comment='#', header=None, delim_whitespace=True)

        xydat = dat[[xcol_, ycol_]]
        zdat, wdat = None, None
        if zcol is not None:
            zdat = dat[zcol_]
        if wcol is not None:
            wdat = dat[wcol_]
        return xydat, zdat, wdat
    else:
        raise NotImplementedError()

def plotter(indata, xp_expr='x', yp_expr='y', xerr_expr=None, yerr_expr=None, label='', rmline=False, rmpoint=False, intpalg=0):
    assert not (rmline and rmpoint) # can not remove the points and lines at the same time
    global imarker, ils, icolor
    x, y = indata[0].values[:, 0], indata[0].values[:, 1]
    valid = np.isfinite(x)&np.isfinite(y)
    if indata[1] is not None:
        has_z = True
        z = indata[1].values
    else:
        has_z = False

    if indata[2] is not None:
        has_w = True
        w = indata[2].values
    else:
        has_w = False

    if not valid.all():
        print('  {}/{} invalid points are removed'.format(valid.size-valid.sum(), valid.size))

    if valid.any():
        x, y = x[valid], y[valid]
        if has_z:
            z = z[valid]
        if has_w:
            w = w[valid]
        xp = eval(xp_expr)
        yp = eval(yp_expr)

        # determine the line style
        if rmline is None:
            myls = 'None'
        else:
            myls = linestyles[ils%nls]
            ils += 1

        #print(xp, yp)
        is_set_label = False
        if yerr_expr or xerr_expr:
            if rmpoint:
                warnings.warn('You can not remove the markers when plotting errorbars', RuntimeWarning)

            xperr, yperr = None, None
            if xerr_expr is not None:
                xperr = eval(xerr_expr)
            if yerr_expr is not None:
                yperr = eval(yerr_expr)

            ax.errorbar(xp, yp, yerr=yperr, xerr=xperr, ls='None',
                        label=label, marker=markers[imarker%nmarker], ms=3, color='C%d'%(icolor%10))
            imarker += 1
            is_set_label = True
        else:
            if not rmpoint:
                marker_kwds = {'marker': markers[imarker%nmarker], 'ms': 3}
                imarker += 1
                ax.plot(xp, yp, ls='None', label=label, c='C%d'%(icolor%10), **marker_kwds)
                is_set_label = True

        if myls != 'None':
            if is_set_label:
                label0 = None
            else:
                label0 = label
            if intpalg==0: # linear
                ax.plot(xp, yp, ls=myls, c='C%d'%(icolor%10), label=label0)
            elif intpalg==1: # spline
                splf = interpolate.CubicSpline(xp, yp)
                xp1 = np.linspace(xp[0], xp[-1], 3000)
                yp1 = splf(xp1)
                ax.plot(xp1, yp1, ls=myls, c='C%d'%(icolor%10), label=label0)
            else:
                raise NotImplementedError()

        icolor += 1
    else:
        print('  no valid points')

def saver(outfile='figure.pdf',
          xmin=None, xmax=None, ymin=None, ymax=None,
          logx=False, logy=False, xlabel='x', ylabel='y',
          grid=False, legend=True):
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_xscale('log' if logx else 'linear')
    ax.set_yscale('log' if logy else 'linear')
    ax.set_xlim(xmin=xmin, xmax=xmax)
    ax.set_ylim(ymin=ymin, ymax=ymax)
    if grid:
        ax.grid(True, ls='--')
    if legend:
        ax.legend()

    if outfile is None:
        print('showing the figure ...')
        plt.show()
    else:
        print('save to {} ...'.format(outfile))
        fg.savefig(outfile)


if __name__ == '__main__':
    import argparse
    import matplotlib as mpl

    parser = argparse.ArgumentParser(prefix_chars='-+',
             description='A quick and dirty way to make a plot')

    parser_io = parser.add_argument_group('IO operations')
    parser_io.add_argument("infile", nargs='+', help='file_name[:label_name] of the input file')
    parser_io.add_argument("-intype", type=str, default='table', choices=['table'], help='the type of input file (default: table)')
    parser_io.add_argument("-outfile", type=str, default=None, help='the name of output figure (default: None)')

    parser_xy = parser.add_argument_group('xy axis')
    parser_xy.add_argument("-xcol", type=str, default='0', help='the col index of x variable (default: 0)')
    parser_xy.add_argument("-ycol", type=str, default='1', help='the col index of y varibale (default: 1)')
    parser_xy.add_argument("-zcol", type=str, default=None, help='the col index of z varibale (default: None)')
    parser_xy.add_argument("-wcol", type=str, default=None, help='the col index of w varibale (default: None)')
    parser_xy.add_argument("-xpexpr", type=str, default='x', help='the expr used as Xp axis in the figure (default: x)')
    parser_xy.add_argument("-ypexpr", type=str, default='y', help='the expr used as Yp axis in the figure (default: y)')
    parser_xy.add_argument("-xperrexpr", type=str, default=None, help='the expr used as Xp error axis in the figure (default: None)')
    parser_xy.add_argument("-yperrexpr", type=str, default=None, help='the expr used as Yp error axis in the figure (default: None)')

    parser_xyp = parser.add_argument_group('xy plot settings')
    parser_xyp.add_argument("-xpmin", type=float, default=None, help='the Xp lower bound in the plot (default: None)')
    parser_xyp.add_argument("-xpmax", type=float, default=None, help='the Xp upper bound in the plot (default: None)')
    parser_xyp.add_argument("-ypmin", type=float, default=None, help='the Yp lower bound in the plot (default: None)')
    parser_xyp.add_argument("-ypmax", type=float, default=None, help='the Yp upper bound in the plot (default: None)')
    parser_xyp.add_argument("+logxp", action='store_true', default=False, help='if the Xp in logscale (default: False)')
    parser_xyp.add_argument("+logyp", action='store_true', default=False, help='if the Yp in logscale (default: False)')
    parser_xyp.add_argument("-xplabel", type=str, default=None, help='the name of Xp axis in the plot (default: None)')
    parser_xyp.add_argument("-yplabel", type=str, default=None, help='the name of Yp axis in the plot  (default: None)')

    parser_msc = parser.add_argument_group('misc plot settings')
    parser_msc.add_argument("-line", action='store_true', default=False, help='if remove the line connecting the points (default: False)')
    parser_msc.add_argument("-points", action='store_true', default=False, help='if remove the markers of the points (default: False)')
    parser_msc.add_argument("-legend", action='store_true', default=False, help='if remove the legend (default: False)')
    parser_msc.add_argument("+grid", action='store_true', default=False, help='if add grids (default: False)')
    parser_msc.add_argument("+spline", action='store_true', default=False, help='if use the spline line to connect points (default: False)')

    args = parser.parse_args()

    if args.outfile is not None:
        mpl.use('Agg')

    from matplotlib import pyplot as plt

    fg, ax = plt.subplots()
    for idx, infile in enumerate(args.infile):
        if ':' in infile:
            infile, label = infile.split(':')
        else:
            infile_simple = infile.split('/')[-1]
            label = '{}[{}]'.format(infile_simple, idx)

        try:
            pdat = loader(infile, type_=args.intype, xcol=args.xcol, ycol=args.ycol, zcol=args.zcol, wcol=args.wcol)
        except Exception as e:
            print('[error] fail to load {}, because of {}'.format(infile, str(e).strip()))
        else:
            print('plotting {} [{}] ...'.format(infile, idx))
            plotter(pdat, xp_expr=args.xpexpr, yp_expr=args.ypexpr, xerr_expr=args.xperrexpr, yerr_expr=args.yperrexpr,
                    label=label, rmline=args.line, rmpoint=args.points, intpalg=1 if args.spline else 0)

    xplabel = 'Xp: {}'.format(args.xpexpr) if args.xplabel is None else args.xplabel
    yplabel = 'Yp: {}'.format(args.ypexpr) if args.yplabel is None else args.yplabel

    if args.xplabel is None or args.yplabel is None:
        annot = 'y: col={}\nx: col={}'.format(ycol_, xcol_)
        if zcol_ is not None:
            annot = 'z: col={}\n'.format(zcol_) + annot
        if wcol_ is not None:
            annot = 'w: col={}\n'.format(wcol_) + annot
        ax.annotate(annot, xy=(10., 10.), xycoords='figure pixels')

    saver(outfile=args.outfile,
          xmin=args.xpmin, xmax=args.xpmax, ymin=args.ypmin, ymax=args.ypmax,
          logx=args.logxp, logy=args.logyp, xlabel=xplabel, ylabel=yplabel,
          grid=args.grid, legend=not args.legend)
